import torch
import torch.nn.functional as F
from tqdm import tqdm
import torch.nn as nn
from dice_loss import dice_coeff

def eval_net(net, loader, device, out_channel=1):
    """Evaluation without the densecrf with the dice coefficient"""
    net.eval()
    mask_type = torch.float32 if out_channel == 1 else torch.long
    n_val = len(loader)  # the number of batch

    tot = 0
  
    with tqdm(total=n_val, desc='Validation round', unit='batch', leave=False) as pbar:
        for imgs, true_masks in loader:

            batch, channe, W, H = imgs.shape
            # imgs = imgs.view((batch, 1, 1, W, H))
            imgs = imgs.to(device=device, dtype=torch.float32)
            true_masks = true_masks.to(device=device, dtype=mask_type)

            with torch.no_grad():
                mask_pred = net(imgs)
                # mask_pred = mask_pred.view((batch, W, H))
                mask_pred = mask_pred.squeeze(1) # 去掉通道那个维度
             
                if true_masks.sum() != 0:
                    tot += dice_coeff(mask_pred, true_masks).item()
                else :
                    n_val -= 1

            pbar.update()

    net.train()
    dice_loss = tot / n_val

    return dice_loss  
